import onnxruntime as ort
import cv2
import numpy as np
class LangCls:
    def __init__(self,ckpt_path,wh=(256,64)):
        self.ort_session = ort.InferenceSession(ckpt_path,providers=['CUDAExecutionProvider'])
        self.wh=wh
        self.lans=['Korean','Hindi','Latin','Chinese']
    def __call__(self, imgs):
        all_imgs=[]
        for img in imgs:
            img=cv2.resize(img,self.wh)
            img=img.astype(np.float32)
            img=(img-128)/128
            img=img[:,:,[2,1,0]]
            all_imgs.append(img)
        all_imgs=np.stack(all_imgs,axis=0)
        all_imgs=np.transpose(all_imgs,axes=(0,3,1,2))
        #print('process time',time()-t1)
        res = self.ort_session.run(["prob"], {'input': all_imgs})[0]
        index=np.argmax(res,axis=-1)
        langs=[]
        for i in index:
            langs.append(self.lans[i])
        return langs

if __name__=='__main__':
    model=LangCls('/media/wsl/a9f0161f-7971-c843-8c81-c68049a0235a/mlt_ckpt/rec/lang_cls.onnx')
    im1=cv2.imread('/media/wsl/a9f0161f-7971-c843-8c81-c68049a0235a/PublicDataSet/OCR/语种判别MLT19/test/Hindi/15.png')
    im2=cv2.imread('/media/wsl/a9f0161f-7971-c843-8c81-c68049a0235a/PublicDataSet/OCR/语种判别MLT19/test/Korean/14.png')
    im3=cv2.imread('/media/wsl/a9f0161f-7971-c843-8c81-c68049a0235a/PublicDataSet/OCR/语种判别MLT19/test/Latin/47.png')
    index=model([im1,im2,im3])
    print(index)